import networks
import argparse
import utils
import torch
import itertools
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from datasets import ImageDataset
import time
import numpy as np



# Get training options from the command line
def get_opt():
    #解析命令行参数，为训练过程提供各种可配置的选项。
    parser = argparse.ArgumentParser()
    '''
    创建一个命令行参数解析器对象。
    argparse 是 Python 标准库中的一个模块，用于从命令行读取参数和选项。
    ArgumentParser() 是 argparse 模块中的一个类，用于创建一个参数解析器对象 parser。
    通过 parser,可以定义程序可以接受哪些命令行参数（比如 --epoch、--batch_size 等），以及这些参数的类型、默认值和帮助信息。
    这样，用户在运行脚本时可以通过命令行传递参数，灵活控制程序的行为，而不需要修改代码。
    '''
    # Training parameters
    parser.add_argument('--epoch', type=int, default=0, help='starting epoch')
    #训练从第几轮开始（通常用于断点续训）。
    parser.add_argument('--n_epochs', type = int, default = 100, help = 'number of epochs with initial learning rate')
    #初始学习率下训练多少轮。
    parser.add_argument('--n_epochs_decay', type = int, default = 100, help = 'number of epochs starting the decay of learning rate')
    #学习率衰减阶段的轮数。
    parser.add_argument('--beta1', type = float, default = 0.5, help = 'momentum term of the Adam optimizer')
    '''
    添加一个命令行参数 --beta1,用于设置 Adam 优化器中的动量参数 β₁(beta1)。
    在深度学习优化中,动量(Momentum)的思想是：当前的参数更新不仅仅依赖于当前的梯度，还会参考之前梯度的累积趋势。
    这样可以让参数在“惯性”的带动下更快地沿着正确的方向前进，减少震荡，加速收敛。
    β₁ 的作用:
    1. 平滑梯度：β₁ 越大，历史梯度的影响越大，更新更平滑，能有效减少训练过程中的震荡。
    2. 加速收敛：在梯度方向一致时，动量会累积，使参数更新更快。
    3. 防止陷入局部最优：动量可以帮助模型跳出一些小的局部最优点。
    4. 调节灵敏度：β₁ 越小，模型对当前梯度更敏感，容易受噪声影响；
    通常 β₁ 取 0.9(论文推荐）,但在生成对抗网络(GAN)等任务中，常用 0.5,可以让训练更稳定。
    调整方法：先用推荐值训练，观察 loss 曲线和平稳性。如果 loss 曲线抖动大，尝试增大 β₁，反之亦然。每次调整幅度建议 0.1~0.2。
    '''
    parser.add_argument('--lr', type = float, default = 0.0002, help = 'initial learning rate')
    '''
    学习率(Learning Rate,通常用 lr 表示）是深度学习和机器学习中用于控制模型参数更新步伐的一个超参数。
    在神经网络训练过程中，模型会通过反向传播算法计算损失函数关于参数的梯度，然后用这些梯度来更新参数，使损失函数逐步减小。
    学习率决定了每次参数更新的“步长”有多大。
    学习率大：每次更新步子大，收敛速度快，但可能会错过最优解，甚至导致训练不稳定。
    学习率小：每次更新步子小，训练更稳定，但收敛速度慢，可能陷入局部最优。
    '''
    parser.add_argument('--batch_size', type = int, default = 1, help = 'batch size of training')
    #每个批次的样本数。
    parser.add_argument('--cuda', action='store_true', help='use GPU computation')
    #是否使用GPU。
    parser.add_argument('--rootdir', type=str, default='datasets/NIRI_to_NIRII/', help='root directory of the dataset')
    #数据集的根目录
    parser.add_argument('--n_cpu', type=int, default=4, help='number of cpu threads to use during batch generation')
    #加载数据时使用的CPU线程数，加快数据预处理。
    parser.add_argument('--u_net', action='store_true', help='use U-net generator')
    #是否使用U-net结构的生成器。加上该参数则用U-net，否则用默认生成器。
    parser.add_argument('--pretrained', action='store_true', help='load pretrained weights')
    #是否加载预训练权重。加上该参数则加载预训练模型。

    # Model parameters，模型参数部分
    parser.add_argument('--sizeh', type=int, default=512, help='size of the image')
    parser.add_argument('--sizew', type=int, default=640, help='size of the image')
    parser.add_argument('--input_nc', type = int, default = 1, help = 'number of input channels')
    #输入图片的通道数。
    parser.add_argument('--output_nc', type = int, default = 1, help = 'number of output channels')
    #输出图片的通道数。
    parser.add_argument('--ngf', type = int, default = 64, help = 'number of filters in the generator')
    #生成器第一层卷积核的数量，影响网络容量。
    parser.add_argument('--ndf', type = int, default = 64, help = 'number of filters in the discriminator')
    #判别器第一层卷积核的数量。
    parser.add_argument('--dropout', type = bool, default = False, help = 'whether to use dropout')
    #是否在网络中使用dropout（防止过拟合）。
    parser.add_argument('--n_res', type = int, default = 9, help = 'number of resNet blocks')
    #生成器中ResNet残差块的数量。
    parser.add_argument('--cycle_loss', type = float, default=10, help = 'coefficient of cycle consistent loss')
    #循环一致性损失的系数（CycleGAN特有）。
    '''
    循环一致性损失(Cycle Consistency Loss)是 CycleGAN 中的一种损失函数，用于鼓励生成器在生成图像时保持输入图像的某些特征。
    它决定了模型在训练时对“循环一致性”约束的重视程度。
    设置得太小，模型可能会忽略循环一致性，导致生成的图像失真或内容丢失；设置得太大，则可能抑制风格迁移的能力，使生成图像过于保守。
    一般设置为 10 或 20 之间的数，具体取值要根据实际任务和实验效果调整。
    CycleGAN 原论文和大多数后续工作都将其设置为 10,这是一个经验上效果较好的默认值。
    如果生成的图像内容丢失严重、结构不保留，可以适当增大（如 15、20)。
    如果你发现风格迁移不明显、生成图像变化太小，可以适当减小（如 5)。
    一般建议在 5~20 之间尝试，极少数任务会用到更大或更小的值。
    如果identity_loss 也启用（非 0)，可以适当减小 cycle_loss,避免总损失过大导致训练不稳定。
    '''
    parser.add_argument('--identity_loss', type = float, default=0, help = 'coefficient of identity loss')
    #身份损失的系数（有时用于提升生成质量）。
    '''
    身份损失(Identity Loss)是 CycleGAN 中的一种损失函数，用于鼓励生成器在生成图像时保持输入图像的某些特征。
    一般设置为 0 或 0.5~10 之间的小数，具体取值要根据实际任务和实验效果调整。
    如果只做普通的风格迁移，身份损失可以不加(即系数为 0)。
    如果希望生成器在输入本域图片时尽量保持原样(比如色彩迁移、去风格化等任务)，可以加上身份损失，常见设置有：
        0.5(CycleGAN 原论文推荐值之一)
        5 或 10(有些任务如色彩迁移、照片到素描等，论文和实践中会用较大系数)
    调参建议:如果发现风格迁移过强、内容丢失，可以适当增大；如果发现迁移不明显，可以适当减小。
    '''

    opt = parser.parse_args()#parse_args() 方法会自动读取你在命令行输入的参数，并根据 add_argument 定义的参数规则进行解析。
    return opt


def main():
    # Get training options，获取训练选项
    opt = get_opt()

    # Define the networks，定义网络
    # netG_A: used to transfer image from domain A to domain B
    # netG_B: used to transfer image from domain B to domain A
    netG_A = networks.Generator(opt.input_nc, opt.output_nc, opt.ngf, opt.n_res, opt.dropout)# 将输入图片从A域转换到B域
    netG_B = networks.Generator(opt.output_nc, opt.input_nc, opt.ngf, opt.n_res, opt.dropout)# 将输入图片从B域转换到A域
    if opt.u_net:# 如果指定 --u_net，则用U-Net结构替代默认生成器。
        netG_A = networks.U_net(opt.input_nc, opt.output_nc, opt.ngf)
        netG_B = networks.U_net(opt.output_nc, opt.input_nc, opt.ngf)

    # netD_A: used to test whether an image is from domain B
    # netD_B: used to test whether an image is from domain A
    netD_A = networks.Discriminator(opt.input_nc, opt.ndf)#判断输入图片是否来自B域
    netD_B = networks.Discriminator(opt.output_nc, opt.ndf)#判断输入图片是否来自A域

    # Initialize the networks
    if opt.cuda:# 如果指定 --cuda，则将网络模型移动到GPU上。
        netG_A.cuda()
        netG_B.cuda()
        netD_A.cuda()
        netD_B.cuda()
    utils.init_weight(netG_A)
    utils.init_weight(netG_B)
    utils.init_weight(netD_A)
    utils.init_weight(netD_B)
    # 初始化网络权重

    if opt.pretrained:#如果指定 --pretrained，则加载预训练模型权重。
        netG_A.load_state_dict(torch.load('pretrained/netG_A.pth'))
        netG_B.load_state_dict(torch.load('pretrained/netG_B.pth'))
        netD_A.load_state_dict(torch.load('pretrained/netD_A.pth'))
        netD_B.load_state_dict(torch.load('pretrained/netD_B.pth'))


    # Define the loss functions，定义损失函数
    criterion_GAN = utils.GANLoss()#对抗损失（判别器真假判别）。
    if opt.cuda:
        criterion_GAN.cuda()

    criterion_cycle = torch.nn.L1Loss()#循环一致性损失（CycleGAN核心）。
    #criterion_cycle = torch.nn.MSELoss()
     # Alternatively, can try MSE cycle consistency loss
     '''
     L1损失函数:衡量两个张量之间的绝对差异，为两个张量之间的差的绝对值误差之和。
     MSE损失函数:衡量两个张量之间的平方差异，为两个张量之间的差的均方误差之和。
     L1Loss 对异常值不敏感,MSELoss 对异常值更敏感，梯度更大，收敛速度可能更快，但也更容易受噪声影响。
     在CycleGAN中,使用L1损失函数,因为L1损失函数对异常值不敏感,能够更好地处理图像中的噪声和细节。
     '''
    criterion_identity = torch.nn.L1Loss()#身份损失（CycleGAN特有）。

    # Define the optimizers,为生成器和判别器分别定义Adam优化器。
    optimizer_G = torch.optim.Adam(itertools.chain(netG_A.parameters(), netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
    # itertools.chain把两个生成器的参数合并成一个大的参数列表，这样可以用一个优化器同时优化两个生成器。
    optimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
    optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
    '''
    torch.optim.Adam是PyTorch中实现的Adam优化器(Adaptive Moment Estimation)，是一种常用的深度学习参数优化算法。
    优化器的作用是根据损失函数的梯度，自动调整神经网络中参数的值，使得损失函数逐步减小，模型性能提升。
    Adam优化器结合了动量(Momentum)和自适应学习率(Adaptive Learning Rate)的思想，能够自适应地调整学习率，在训练过程中更加稳定和高效。
    常用参数：
    lr: 学习率，控制参数更新的步长。
    betas: 动量参数，控制梯度下降的更新速度。
    eps: 防止除零的小常数。
    weight_decay: 正则化项，防止过拟合。
    amsgrad: 是否使用AMSGrad变体,提高收敛稳定性。
    '''

    # Create learning rate schedulers，定义学习率调度器，实现学习率的线性衰减。
    lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda = utils.Lambda_rule(opt.epoch, opt.n_epochs, opt.n_epochs_decay).step)
    lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda = utils.Lambda_rule(opt.epoch, opt.n_epochs, opt.n_epochs_decay).step)
    lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda = utils.Lambda_rule(opt.epoch, opt.n_epochs, opt.n_epochs_decay).step)
    '''
    torch.optim.lr_scheduler.LambdaLR
    LambdaLR是PyTorch中实现的学习率调度器用于在训练过程中动态调整学习率。
    LambdaLR允许你自定义一个函数(lambda),这个函数根据当前的epoch返回一个缩放因子,动态调整学习率。
    每次调用scheduler.step()时,优化器的学习率会被乘以这个缩放因子。
    这个函数可以实现学习率的线性衰减、指数衰减、余弦衰减等不同的衰减方式。
    常用参数：
    optimizer: 要调整学习率的优化器。
    lr_lambda: 一个lambda函数,用于计算新的学习率。
    last_epoch: 上一个epoch的索引,用于确定当前epoch的索引。
    verbose: 是否打印学习率调度信息。
    '''

    # 定义两个张量，用于存储输入图片。
    Tensor = torch.cuda.FloatTensor if opt.cuda else torch.Tensor
    input_A = Tensor(opt.batch_size, opt.input_nc, opt.sizeh, opt.sizew)
    input_B = Tensor(opt.batch_size, opt.output_nc, opt.sizeh, opt.sizew)
    # 根据是否使用GPU，选择合适的张量类型，并为A域和B域的输入图片分配内存空间，方便后续的数据加载和模型训练。

    # Define two image pools to store generated images，定义两个图片池，用于存储生成的图片。
    fake_A_pool = utils.ImagePool()
    fake_B_pool = utils.ImagePool()
    # ImagePool类的作用是：存储和重复使用生成的假图片，避免每次训练都重新生成，节省计算资源。
    # 当生成器生成新的假图片时，会先检查图片池中是否已经有类似的图片，如果有，则直接使用池中的图片，而不是重新生成。
    # 这样可以减少计算量，提高训练效率。

    # Define the transform, and load the data
    # 定义一个图像预处理的变换流程，用于在加载训练数据时对每张图片进行一系列标准化和增强操作。
    transform = transforms.Compose([transforms.Resize((opt.sizeh, opt.sizew)),#将图片缩放到指定大小。
                'transforms.Compose 会把括号里的所有变换按顺序组合起来，依次作用于每一张图片。'
                transforms.RandomHorizontalFlip(),#随机水平翻转。
                transforms.ToTensor(),#将PIL图像转换为张量。
                transforms.Normalize((0.5,), (0.5,))])#标准化处理，将像素值缩放到[-1,1]之间。
    dataloader = DataLoader(ImageDataset(opt.rootdir, transform = transform, mode = 'train'), batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu)
    # 定义一个数据加载器，用于加载训练数据。
    # ImageDataset是自定义的数据集类，用于加载和处理训练数据。
    # transform参数用于指定对每张图片进行的预处理操作。
    # mode参数用于指定数据集的类型，这里设置为'train'，表示加载训练数据。
    # batch_size参数用于指定每个批次的样本数。
    # shuffle参数用于指定是否打乱数据集。
    # num_workers参数用于指定加载数据时使用的CPU线程数。

    # numpy arrays to store the loss of epoch，损失数组，用于存储每个epoch的损失值。
    loss_G_array = np.zeros(opt.n_epochs + opt.n_epochs_decay)
    loss_D_A_array = np.zeros(opt.n_epochs + opt.n_epochs_decay)
    loss_D_B_array = np.zeros(opt.n_epochs + opt.n_epochs_decay)
    # 这三行代码的本质作用是为后续记录和分析训练过程中的损失变化做准备。
    # 每训练完一个 epoch，就会把该 epoch 的损失值存入对应的数组位置，最终可以用这些数组画出损失曲线，分析模型训练效果和收敛情况。

    # Training，训练过程。
    for epoch in range(opt.epoch, opt.n_epochs + opt.n_epochs_decay):#开始训练主循环，遍历每一个 epoch（训练轮次）。
        start = time.strftime("%H:%M:%S")#获取当前时间。
        print("current epoch :", epoch, " start time :", start)#打印当前epoch的开始时间。
        # Empty list to store the loss of each mini-batch，空列表，用于存储每个mini-batch的损失值。
        loss_G_list = []#生成器损失列表。
        loss_D_A_list = []#判别器A损失列表。
        loss_D_B_list = []#判别器B损失列表。

        for i, batch in enumerate(dataloader):
            if i % 50 == 1:
                print("current step: ", i)
                current = time.strftime("%H:%M:%S")
                print("current time :", current)
                print("last loss G:", loss_G_list[-1], "last loss D_A", loss_D_A_list[-1], "last loss D_B", loss_D_B_list[-1])
                #-1表示列表的最后一个元素，即最近一次的损失值。
            real_A = input_A.copy_(batch['A'])#将当前批次的A域图片复制到input_A张量中。
            real_B = input_B.copy_(batch['B'])#将当前批次的B域图片复制到input_B张量中。
            #每隔 50 个 batch 会打印一次当前进度、时间和最近的损失，方便监控训练过程。

            # Train the generator
            optimizer_G.zero_grad()#清空生成器的梯度。
            '如果不清零，梯度会在每个 batch 之间累加，导致参数更新不是基于当前 batch 的梯度，而是所有 batch 的累积梯度，训练结果会异常。'
            '正确做法是：每个 batch 开始前都要清零梯度，这样每次参数更新只基于当前 batch 的梯度。'

            # Compute fake images and reconstructed images
            fake_B = netG_A(real_A)
            fake_A = netG_B(real_B)

            if opt.identity_loss != 0:
                same_B = netG_A(real_B)
                same_A = netG_B(real_A)

            # discriminators require no gradients when optimizing generators,判别器在优化生成器时不需要梯度。
            utils.set_requires_grad([netD_A, netD_B], False)

            # Identity loss
            if opt.identity_loss != 0:
                loss_identity_A = criterion_identity(same_A, real_A) * opt.identity_loss
                loss_identity_B = criterion_identity(same_B, real_B) * opt.identity_loss

            # GAN loss
            prediction_fake_B = netD_B(fake_B)
            loss_gan_B = criterion_GAN(prediction_fake_B, True)
            prediction_fake_A = netD_A(fake_A)
            loss_gan_A = criterion_GAN(prediction_fake_A, True)

            # Cycle consistent loss
            recA = netG_B(fake_B)
            recB = netG_A(fake_A)
            loss_cycle_A = criterion_cycle(recA, real_A) * opt.cycle_loss
            loss_cycle_B = criterion_cycle(recB, real_B) * opt.cycle_loss

            # total loss without the identity loss
            loss_G = loss_gan_B + loss_gan_A + loss_cycle_A + loss_cycle_B

            if opt.identity_loss != 0:
                loss_G += loss_identity_A + loss_identity_B

            loss_G_list.append(loss_G.item())
            'loss_G 是一个只有一个元素的张量(比如 tensor(2.345))，而不是普通的 Python 数值类型(如 float)。'
            '.item() 是 PyTorch 张量的方法，用于将只包含一个元素的张量转换为对应的 Python 数值(float 或 int)'
            '如果直接 append(loss_G)，会把张量对象加到列表里，后续做数值运算(如求平均、画图)会不方便。'
            '用 .item() 转成 float,列表里全是数值，后续处理更方便。'
            loss_G.backward()
            optimizer_G.step()

            # Train the discriminator
            utils.set_requires_grad([netD_A, netD_B], True)


            # Train the discriminator D_A
            optimizer_D_A.zero_grad()
            # real images
            pred_real = netD_A(real_A)
            loss_D_real = criterion_GAN(pred_real, True)

            # fake images
            fake_A = fake_A_pool.query(fake_A)
            pred_fake = netD_A(fake_A.detach())
            loss_D_fake = criterion_GAN(pred_fake, False)

            #total loss
            loss_D_A = (loss_D_real + loss_D_fake) * 0.5
            loss_D_A_list.append(loss_D_A.item())
            loss_D_A.backward()
            optimizer_D_A.step()

            # Train the discriminator D_B
            optimizer_D_B.zero_grad()
            # real images
            pred_real = netD_B(real_B)
            loss_D_real = criterion_GAN(pred_real, True)

            # fake images
            fake_B = fake_B_pool.query(fake_B)
            pred_fake = netD_B(fake_B.detach())
            loss_D_fake = criterion_GAN(pred_fake, False)

            # total loss
            loss_D_B = (loss_D_real + loss_D_fake) * 0.5
            loss_D_B_list.append(loss_D_B.item())
            loss_D_B.backward()
            optimizer_D_B.step()
        # 生成器损失 = 对抗损失（骗判别器）+ 循环一致性损失（保证可逆）+（可选）身份损失（防止过度风格化）。
        # 判别器损失 = 区分真实和生成图片的能力。

        # Update the learning rate
        lr_scheduler_G.step()
        lr_scheduler_D_A.step()
        lr_scheduler_D_B.step()

        # Save models checkpoints，保存模型检查点。
        torch.save(netG_A.state_dict(), 'model/netG_A.pth')
        torch.save(netG_B.state_dict(), 'model/netG_B.pth')
        torch.save(netD_A.state_dict(), 'model/netD_A.pth')
        torch.save(netD_B.state_dict(), 'model/netD_B.pth')




        # Save other checkpoint information，保存其他检查点信息。
        checkpoint = {'epoch': epoch,
                      'optimizer_G': optimizer_G.state_dict(),
                      'optimizer_D_A': optimizer_D_A.state_dict(),
                      'optimizer_D_B': optimizer_D_B.state_dict(),
                      'lr_scheduler_G': lr_scheduler_G.state_dict(),
                      'lr_scheduler_D_A': lr_scheduler_D_A.state_dict(),
                      'lr_scheduler_D_B': lr_scheduler_D_B.state_dict()}
        torch.save(checkpoint, 'model/checkpoint.pth')



        # Update the numpy arrays that record the loss，更新损失数组。
        loss_G_array[epoch] = sum(loss_G_list) / len(loss_G_list)
        loss_D_A_array[epoch] = sum(loss_D_A_list) / len(loss_D_A_list)
        loss_D_B_array[epoch] = sum(loss_D_B_list) / len(loss_D_B_list)
        np.savetxt('model/loss_G.txt', loss_G_array)
        np.savetxt('model/loss_D_A.txt', loss_D_A_array)
        np.savetxt('model/loss_D_b.txt', loss_D_B_array)


        if epoch % 10 == 9:#每 10 个 epoch 额外保存一次模型快照，方便调试和恢复。
            torch.save(netG_A.state_dict(), 'model/netG_A' + str(epoch) + '.pth')
            torch.save(netG_B.state_dict(), 'model/netG_B' + str(epoch) + '.pth')
            torch.save(netD_A.state_dict(), 'model/netD_A' + str(epoch) + '.pth')
            torch.save(netD_B.state_dict(), 'model/netD_B' + str(epoch) + '.pth')

        end = time.strftime("%H:%M:%S")
        print("current epoch :", epoch, " end time :", end)
        print("G loss :", loss_G_array[epoch], "D_A loss :", loss_D_A_array[epoch], "D_B loss :", loss_D_B_array[epoch])



if __name__ == "__main__":
    main()

